import json
import traceback

import gym
import random
import requests
import string
import time
import re
from openai import OpenAI


class MirrirAPIAgentSiteEnv(object):
    """Gym environment for HTML mode of WebShop environment"""

    def __init__(self,max_steps, **kwargs):
        """
        Constructor for HTML environment

        Arguments:
        observation_mode (`str`) -- ['html' | 'text'] (default 'html')
        pause (`float`) -- Pause (in seconds) after taking an action.
            This is mainly for demo purposes.
            Recommended value: 2.0s
        render (`bool`) -- Show browser if set to `True`.
        session ('str') -- Session ID to initialize environment with
        """
        # self.reset()
        self.max_step = max_steps
        self._step = 0
        self.qwen_urls = [
            'http://10.3.14.10:8001/v1/chat/completions',
        ]
        # 'http://10.3.14.10:1091/v1/chat/completions',
        # 'http://10.3.14.10:1101/v1/chat/completions',
        # 'http://10.3.14.10:1111/v1/chat/completions',


        self.observation = None
        self.query = None
        # self.rewardClient = OpenAI(
        #     # api_key="sk-1a4738bb16594d2eb729dc8c33ea5044",
        #     api_key="sk-301d7a25146a45f9affcb4b8f4944979",
        #     base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        # )
        # self.rewardClient = OpenAI(
        #     api_key="2SXU7yKBLAZJ0TUOAyUJeLPwNVeqvGskCjsylrmuhYnzsBMa4BWQzZszF9V_dFrf4hNEAVA1ksvUTe-ZZOkv3Q",
        #     base_url="https://api.modelarts-maas.com/v1",
        #     timeout=300,
        #     max_retries=10,
        # )
        self.rewardClient = OpenAI(
            api_key="XXXXFAKE_KEY_DO_NOT_REMOVE",
            base_url="http://10.3.14.10:1024/v1",
            timeout=300,
            max_retries=10,
        )
        # Set flags and values for WebShop session
        self.assigned_session = kwargs.get('session')
        self.session = None
        with open("/mnt/nas/tianye/code/agenticsystem/agent_system/environments/env_package/mirrorapi/mirrirapi/data/G3_query.json", "r") as f:
            self.datasets = json.load(f)
        self.index = None
        with open("/mnt/nas/tianye/code/agenticsystem/agent_system/environments/env_package/mirrorapi/mirrirapi/prompts/mirrorapi_system_prompt.txt", 'r', encoding='utf-8') as f:
            self.system_prompt = f.read()
        with open("/mnt/nas/tianye/code/agenticsystem/agent_system/environments/env_package/mirrorapi/mirrirapi/prompts/evaluator_system_prompt.txt", 'r', encoding='utf-8') as f:
            self.eval_system_prompt = f.read()
        
        
        random.seed(time.time_ns())
        # self.reset()

    def initialize(self, goals=None):
        """

        """
        # todo
        print("Warnning: do nothing!!!")
        class SimServer():
            def __init__(self):
                self.goals = list(range(503))
        self.server = SimServer()
        self.initialized = True

    def extract_api_feedback(self, api_result) -> str:
        result = json.loads(api_result)
        api_feedback = result.get('choices')[0]["message"]["content"]
        return api_feedback

    def api_calling(self, system_prompt, user_prompt):
        """Strip HTML of tags and add separators to convert observation into simple mode"""
        url = random.choice(self.qwen_urls)
        # url = 'http://10.3.14.10:12345/v1/chat/completions'
        payload = {
            "model": "MirrorAPI",
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 0.7,
            "max_tokens": 8192
        }
        headers = {
            'accept': 'application/json',
            'Content-Type': 'application/json',
            'Authorization': "Bearer EMPTY"
        }
        # print(f"###########  url {url}, header {headers}, data {json.dumps(payload)}")

        response = requests.post(url, headers=headers, data=json.dumps(payload),timeout=300)
        # print(f"### response {response}")
        api_feedback = self.extract_api_feedback(response.text)
        return api_feedback

    def extract_action_info(self, action):
        result = {
            "think": "",
            "action": "",
            "invoke": "",
            "observation": "",
            "answer": "",
        }
        match_think = re.search(r"<think>(.*?)</think>", action, re.DOTALL)
        match_action = re.search(r"<action>(.*?)</action>", action, re.DOTALL)
        match_invoke = re.search(r"<invoke>(.*?)</invoke>", action, re.DOTALL)
        match_observation = re.search(r"<invoke>(.*?)</invoke>", action, re.DOTALL)
        match_answer = re.search(r"<answer>(.*?)</answer>", action, re.DOTALL)
        if match_think:
            result["think"] = match_think.group(1).strip()
        if match_action:
            result["action"] = match_action.group(1).strip()
        if match_invoke:
            result["invoke"] = match_invoke.group(1).strip()
        if match_observation:
            result["observation"] = match_observation.group(1).strip()
        if match_answer:
            result["answer"] = match_answer.group(1).strip()
        return result

    def get_reward(self, query: str, answer: str, num_his: int):
        validater = {
            "query": query,
            "answer": answer
        }
        # return 0
        try:
            response = self.rewardClient.chat.completions.create(
                model="models--Qwen--Qwen2.5-32B-Instruct",
                messages=[
                    {
                        "role": "system",
                        "content": self.eval_system_prompt
                    },
                    {
                        "role": "user",
                        "content": json.dumps(validater)
                    }
                ],
                temperature=0.4,
                n=1,
                # stop=None,
                stream=False,
            )
            action = response.choices[0].message.content.strip()
            # action = "Solved"
            #如果有奖励结果，就不给步骤分数
            reward = 12/(num_his/self.max_step)-num_his if action == "Solved" else -1*num_his
            return reward
        except Exception as e:
            traceback.print_exc()
            return 0


    def step(self, action, invoke, valid, assistant_turns):
        """
        Takes an action, updates WebShop environment, and returns (observation, reward, done, info)

        Arguments:
        action (`str`): An action should be of the following structure:
          - search[keywords]
          - click[value]
        If action not valid, perform nothing.
        """
        # breakpoint()
        observation = " "
        reward = 0.0
        done = False
        info = {}
        info["available_actions"] = self.obtain_api_list(self.index)
        self._step += 1
        print(f" current step {self._step} , max step {self.max_step}")
        done = self._step >= self.max_step
        if valid:
            #print(f"### action {action}")
            # if invoke is None or "api_name" not in invoke:
            #     invoke = json.loads("\n{\n    \"category\": \"Food\",\n    \"tool_name\": \"The Cocktail DB\",\n    \"api_name\": \"List of Cocktails\",\n    \"tool_input\": \"{{}}\",\n    \"strip\": \"\"\n}\n")
            if invoke != "Answer":
                try:
                    invoke = json.loads(invoke)
                except:
                    print(f"@@@ invoke : {invoke}")
                if "api_name" not in invoke:
                    return "error", -1, done, info
                api_name = invoke["api_name"]
                api_struct = None
                for item in self.info["available_actions"]:
                    if item.get("api_name") == api_name:
                        api_struct = item
                        break
                if api_struct is None:
                    return "error", -1, done, info
                history = ""
                # if his_text_ob != " ":
                #     print(f"@@@ his text ob {len(his_text_ob)}")
                #     for index, item in enumerate(his_text_ob):
                #         history += f"\n Historical API Response {index + 1}: " + item + " \n"
                user_prompt = "API Docs: \n" + json.dumps(api_struct) + "\nRequest: \n" + json.dumps(invoke) + history
                print(f"@@@ api call user_prompt {len(user_prompt)}")
                print(f"@@@ api call api_struct {len(json.dumps(api_struct))}")
                print(f"@@@ api call invoke {len(json.dumps(invoke))}")
                print(f"@@@ api call history len {len(history)}")
                print(f"@@@ api call history {history}")
                observation = self.api_calling(self.system_prompt, user_prompt)
                # print(f"### api calling res {observation}")
                # 达到最大步数但未完成任务获得0奖励，否则获得1奖励为过程奖励
                return observation, 0 if done else 1, done, info
            else:
                try:
                    reward = self.get_reward(self.query,action, assistant_turns)
                except Exception as e:
                    traceback.print_exc()
                    # print(repr(e))
                    reward = 0
                    print(f"### reward error {action} , q {self.query}")
                done = True
                return "done", reward, done, info
        else:
            return "action is not valid.", -1, done, info

    def get_available_actions(self):
       return

    def get_instruction_text(self):
        """Get corresponding instruction text for environment current step"""
        html_obj = self._parse_html(self.browser.page_source)
        instruction_text = html_obj.find(id='instruction-text').h4.text
        return instruction_text


    @property
    def action_space(self):
        # Recommended to use `get_available_actions` instead
        return NotImplementedError

    @property
    def observation_space(self):
        return NotImplementedError

    def obtain_query(self, index):
        query = self.datasets[index]["query"]
        return str(query)

    def obtain_api_list(self, index):
        api_list = self.datasets[index]["api_list"]
        return api_list

    def reset(self, task, api_list, index=0):
        #add observation and info as self
        """Create a new session and reset environment variables"""
        # self.query = self.obtain_query(index)
        self.query = task
        # self.observation = self.obtain_api_list(index)
        self.info = {}
        # self.info["available_actions"] = self.obtain_api_list(index)
        self.info["available_actions"] = api_list
        self.index = index
        self._step = 0
        return self.query, self.info

    def render(self, mode='human'):
        # TODO: Render observation in terminal or WebShop website
        return NotImplementedError

    def close(self):
        # TODO: When DB used instead of JSONs, tear down DB here
        pass



def tag_visible(element):
    """Helper method to strip HTML block of extraneous tags"""
    ignore = {'style', 'script', 'head', 'title', 'meta', '[document]'}
    return (
            element.parent.name not in ignore and not isinstance(element, Comment)
    )